"""
    策略1， 根据 日线的close 的 fib 历史数据， 预测 1， 2 ，3 ，5  日的收益
"""
import talib
import pandas as pd
import numpy as np
import talib as ta
import datetime

from fib_config import stime, stime1, etime, etime1
from fib_utils import fib_yield_for, myroll, roll_corr, myroll_apply_corr, corr_v1
stock_list = []


def fib_rate(fib_df, column, n=9):
    for i in fib_yield_for(n):
        fib_df[column + "_fib" + str(i)] = fib_df[column] / fib_df[column].shift(i) - 1.0


data = get_price(stock_list, stime, etime,
                         '1d',
                         ['open', 'high', 'low', 'close', 'volume',
                         #   涨跌幅        振幅         换手率
                          'quote_rate', 'amp_rate', 'turnover_rate'],
                         True,  # 是否跳过停牌
                         "pre",  # 前复权
                         0,  # 天数
                         is_panel=1)
df = data.to_frame().reset_index()
# df.columns = ['date', 'symbol', 'close', 'quote_rate', 'amp_rate', 'turnover_rate']
df.rename(columns={'major': 'date','minor': 'symbol'}, inplace=True)
df = df.sort_values(['symbol', 'date'])
df = df.reset_index(drop=True)


sszs = get_price(["000001.SH"], stime, etime,
                         '1d',
                         #   close      涨跌幅        振幅         换手率
                         ['close'],
                         True,  # 是否跳过停牌
                         "pre",  # 前复权
                         0,  # 天数
                         is_panel=0)["000001.SH"]
sszs.rename(columns={"close": "szss_close"}, inplace=True)

df = pd.merge(df, sszs, left_on="date", right_index=True)
# df.to_csv("./stock_with_szzs.csv")

# 普通因子
# 周线换手率均值, 月线换手率均值
df['close_ma5'] = df.groupby(['symbol'])['close'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_ma20'] = df.groupby(['symbol'])['close'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_ma60'] = df.groupby(['symbol'])['close'].rolling(60).mean().reset_index(drop=True, level=0)
# n天的平均价到今天的涨幅
df["close" + "_fib_d1"] = df["close"] / df["close"].shift(1) - 1.0
df["close" + "_fib_d2"] = df["close"] / df["close"].shift(2) - 1.0
df["close" + "_fib_d3"] = df["close"] / df["close"].shift(3) - 1.0
# n周的平均价到今天的涨幅
df["close" + "_fib_w1"] = df["close"] / df["close_ma5"].shift(1) - 1.0
df["close" + "_fib_w2"] = df["close"] / df["close_ma5"].shift(2) - 1.0
df["close" + "_fib_w3"] = df["close"] / df["close_ma5"].shift(3) - 1.0
# n月的平均价到今天的涨幅
df["close" + "_fib_m1"] = df["close"] / df["close_ma20"].shift(1) - 1.0
df["close" + "_fib_m2"] = df["close"] / df["close_ma20"].shift(2) - 1.0
df["close" + "_fib_m3"] = df["close"] / df["close_ma20"].shift(3) - 1.0

df["close" + "_fib_j1"] = df["close"] / df["close_ma60"] - 1.0

# 周线换手率均值, 月线换手率均值
df['turnover_r5'] = df.groupby(['symbol'])['turnover_rate'].rolling(5).mean().reset_index(drop=True, level=0)
df['turnover_r20'] = df.groupby(['symbol'])['turnover_rate'].rolling(20).mean().reset_index(drop=True, level=0)
df['turnover_r60'] = df.groupby(['symbol'])['turnover_rate'].rolling(60).mean().reset_index(drop=True, level=0)

df["turnover" + "_fib_d1"] = df.groupby(['symbol'])['turnover_rate'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_d2"] = df.groupby(['symbol'])['turnover_rate'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_d3"] = df.groupby(['symbol'])['turnover_rate'].shift(3).reset_index(drop=True, level=0)

df["turnover" + "_fib_w1"] = df.groupby(['symbol'])['turnover_r5'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_w2"] = df.groupby(['symbol'])['turnover_r5'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_w3"] = df.groupby(['symbol'])['turnover_r5'].shift(3).reset_index(drop=True, level=0)

df["turnover" + "_fib_m1"] = df.groupby(['symbol'])['turnover_r20'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_m2"] = df.groupby(['symbol'])['turnover_r20'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_m3"] = df.groupby(['symbol'])['turnover_r20'].shift(3).reset_index(drop=True, level=0)

### 和上证指数相关的因子
# df = df.groupby(["symbol"]).apply(lambda x: myroll_apply_corr(x, 5))
df = df.groupby(['symbol']).apply(corr_v1, "close", "szss_close", [5, 10, 15, 20, 60])
df.sort_index(inplace=True)


# talib ...
def talib_pattern(df1):
    df1['crows2'] = talib.CDL2CROWS(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)/100
    df1['crows3'] = talib.CDL3BLACKCROWS(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)/100
    df1['inside3'] = talib.CDL3INSIDE(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['line_strike'] = talib.CDL3LINESTRIKE(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['outside'] = talib.CDL3OUTSIDE(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['stars_in_south3'] = talib.CDL3STARSINSOUTH(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['wihite_soliders3'] = talib.CDL3WHITESOLDIERS(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    #。。。待续
    return df1

def talib_mom(df1):

    df1['adx_14'] = talib.ADX(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['adxr_14'] = talib.ADXR(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['apo'] = talib.APO(df1['close'].values, fastperiod=12, slowperiod=26, matype=0)
    df1['aroon_up'], df1["aroon_down"] = talib.AROON(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['aroon_osc'] = talib.AROONOSC(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['bop'] = talib.BOP(df1['open'].values, df1['high'].values, df1['low'].values, df1['close'].values)
    df1['cci'] = talib.CCI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['cmo'] = talib.CMO(df1['close'].values, timeperiod=14)

    df1['dx'] = talib.DX(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1["macd"], df1["macd_signal"], df1["macd_hist"] = talib.MACD(df1['close'].values, fastperiod=12, slowperiod=26, signalperiod=9)
    df1['mfi'] = talib.MFI(df1['high'].values, df1['low'].values, df1['close'].values, df1['volume'].values, timeperiod=14)
    df1['di'] = talib.MINUS_DI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['dm'] = talib.MINUS_DM(df1['high'].values, df1['low'].values, timeperiod=14)
    df1['mom'] = talib.MOM(df1['close'].values, timeperiod=10)

    df1['plus_di'] = talib.PLUS_DI(df1['high'].values, df1['low'].values, df1['close'].values, timeperiod=14)
    df1['plus_dm'] = talib.PLUS_DM(df1['high'].values, df1['low'].values, timeperiod=14)
    # 。。。待续
    return df1

def calc_var(stocks):
    # symbol = df1["symbol"].iloc[0]
    data = get_price(stocks, stime, etime,
                     '1m',
                     #   close
                     ['close'],
                     True,  # 是否跳过停牌
                     "pre",  # 前复权
                     0,  # 天数
                     is_panel=1)

    dfm = data.to_frame().reset_index()
    dfm.rename(columns={'major': 'date', 'minor': 'symbol'}, inplace=True)
    dfm["day"] = dfm["date"].apply(lambda x: x.strftime("%Y-%m-%d"))
    dfm = dfm.sort_values(['symbol', 'date'])
    dfm = dfm.reset_index(drop=True)
    df_var = dfm.groupby(['symbol', "day"])[['close']].std()
    df_var = df_var.reset_index()
    df_var.rename(columns={'close': 'close_var'}, inplace=True)
    return df_var

cache_df = None
for stock in stock_list:
    #print("------")
    mdf = pd.merge(df[df['symbol']==stock], calc_var([stock]), on=["symbol", "day"])
    # print(type(mdf.index))
    if cache_df is None:
        cache_df = mdf
    else:
        cache_df = pd.concat([cache_df, mdf])
df = cache_df.reset_index()

df = df.groupby(['symbol']).apply(talib_pattern)
df = df.groupby(['symbol']).apply(talib_mom)

# 收盘价比例， 结果中为NaN的地方表示空值
df['return1'] = df['close'] / df.groupby(['symbol'])['close'].shift(-1) - 1.0
df['return2'] = df['close'] / df.groupby(['symbol'])['close'].shift(-2) - 1.0
df['return3'] = df['close'] / df.groupby(['symbol'])['close'].shift(-3) - 1.0
df['return5'] = df['close'] / df.groupby(['symbol'])['close'].shift(-5) - 1.0

# 删除NAN 的行
df = df.dropna(axis=0, how='any')
# 保留小数
round_map = {
    "quote_rate": 3, "amp_rate": 3, "turnover_rate": 3,
    "close_ma5": 3, "close_ma20": 3, "turnover_r5": 3, "turnover_r20": 3,

    "close_fib_d1": 3, "close_fib_d2": 3, "close_fib_d3": 3,
    "close_fib_w1": 3, "close_fib_w2": 3, "close_fib_w3": 3,
    "close_fib_m1": 3, "close_fib_m2": 3, "close_fib_m3": 3,

    "turnover_fib_d1": 3, "turnover_fib_d2": 3, "turnover_fib_d3": 3,
    "turnover_fib_w1": 3, "turnover_fib_w2": 3, "turnover_fib_w3": 3,
    "turnover_fib_m1": 3, "turnover_fib_m2": 3, "turnover_fib_m3": 3,

    "return1": 3, "return2": 3, "return3": 3, "return5": 3
}
# df = df.round(round_map)
df = df.round(3)

# 删除不需要的列, close 不是比例值
df.drop(['open', 'high', 'low', 'volume'],
        axis=1,
        inplace=True)
# todo 会有 np.inf 出现
df['close_var_ma5'] = df.groupby(['symbol'])['close_var'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_var_ma20'] = df.groupby(['symbol'])['close_var'].rolling(20).mean().reset_index(drop=True, level=0)
df['close_var_ma60'] = df.groupby(['symbol'])['close_var'].rolling(60).mean().reset_index(drop=True, level=0)
# n天的平均价到今天的涨幅
# df["close_var" + "_fib_d1"] = df["close_var"] / df["close_var"].shift(1) - 1.0
# df["close_var" + "_fib_d2"] = df["close_var"] / df["close_var"].shift(2) - 1.0
# df["close_var" + "_fib_d3"] = df["close_var"] / df["close_var"].shift(3) - 1.0
# n周的平均价到今天的涨幅
df["close_var" + "_fib_w1"] = df["close_var"] / df["close_var_ma5"].shift(1) - 1.0
df["close_var" + "_fib_w2"] = df["close_var"] / df["close_var_ma5"].shift(2) - 1.0
df["close_var" + "_fib_w3"] = df["close_var"] / df["close_var_ma5"].shift(3) - 1.0
# n月的平均价到今天的涨幅
df["close_var" + "_fib_m1"] = df["close_var"] / df["close_var_ma20"].shift(1) - 1.0
df["close_var" + "_fib_m2"] = df["close_var"] / df["close_var_ma20"].shift(2) - 1.0
df["close_var" + "_fib_m3"] = df["close_var"] / df["close_var_ma20"].shift(3) - 1.0

def calc_rate(df1):
    df1["turn_rrw1"] = df1["turnover_fib_w1"] / df1["turnover_r60"] - 1
    df1["turn_rrw2"] = df1["turnover_fib_w2"] / df1["turnover_r60"] - 1

    df1["corr_rrw1"] = df1["corr5"] / df1["corr60"] - 1
    df1["corr_rrw2"] = df1["corr10"] / df1["corr60"] - 1
    df1["corr_rrw3"] = df1["corr20"] / df1["corr60"] - 1
    return df1

df = df.groupby(['symbol']).apply(calc_rate)

df.to_csv(r'./fib_close_v1.csv', index=False)

